package yagoMarried;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashSet;

public class ExtractYago {
	
	public ArrayList<YagoEntity> retrieveAll(String relation) throws ClassNotFoundException, SQLException, UnsupportedEncodingException{
		
		ArrayList<YagoEntity>list = new ArrayList<YagoEntity>();
		Class.forName("org.postgresql.Driver");
		Connection db = DriverManager.getConnection("jdbc:postgresql:yago2spotlx_db", "lyao1", "");
		if( db==null){
		//	System.out.println("null");
		}else{
		//	System.out.println("not null");
		}
		Statement st = db.createStatement();
	
		ResultSet rs = st.executeQuery("select arg1,arg2,relation,timebegin,timeend from relationalfacts where relation='"+relation+"' and timebegin is not null limit 10000;");
		
		int find_count =0;
		while (rs.next()) {
		    //out.print("Column 1 returned ");
			String arg1 = rs.getString(1);
			String arg2 = rs.getString(2);
			String rel = rs.getString(3);
			Timestamp timebegin = rs.getTimestamp(4);
			Timestamp timeend = rs.getTimestamp(5);
			YagoEntity ye = new YagoEntity(arg1,arg2,rel,timebegin,timeend);
			list.add(ye);
		   
		}
		rs.close();
		st.close();
		return list;
		
		
		
	}
	
	public ArrayList<YagoEntity> getIndexedEntities(){
		try {
			//File dataFile = new File(dir);
			//BufferedWriter output = new BufferedWriter( new FileWriter(dataFile) );
		
			ArrayList<YagoEntity> plist = this.retrieveAll("isMarriedTo");
			IndexTACReader itacr = new IndexTACReader("/largedata1/cmput696/students/lyao1/tacindex");
			ArrayList<YagoEntity> indexedEntities = new ArrayList<YagoEntity>();
			
			
			
			for(YagoEntity e:plist){
				String p1 = e.getArg1();
				String p2 = itacr.convert(p1);
					String content = itacr.queryWikiContent(p2);
					if(content!=null){

						indexedEntities.add(e);
						
					}else{
						String p3 = itacr.convert(e.getArg2());
						content = itacr.queryWikiContent(p3);
						if( content!=null){
							indexedEntities.add(e);
						}
					}
				
				
			}
			return indexedEntities;
			
			
		} catch (UnsupportedEncodingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (SQLException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		}
		return null;
	}
	//retrieve all playsFor and writes their wiki content into a single file
/*	public void writePlaysFor(String dir){
		try {
			//File dataFile = new File(dir);
			//BufferedWriter output = new BufferedWriter( new FileWriter(dataFile) );
		
			ArrayList<YagoEntity> plist = this.retrieveAll("playsFor");
			IndexTACReader itacr = new IndexTACReader("/largedata1/cmput696/students/lyao1/tacindex");
			
			int found=0;
			HashSet<String> players = new HashSet<String>();
			
			for(YagoEntity e:plist){
				String p1 = e.getArg1();
				
				if( !players.contains(p1)){
					players.add(p1);
					String p2 = itacr.convert(p1);
					
				
					String content = itacr.queryWikiContent(p2);
					if(content!=null){
						
						found++;
						System.out.println("Found:"+p2);
						System.out.println("writing file");
						
						File newf = new File(dir+p2+".txt");
						newf.createNewFile();
						BufferedWriter output = new BufferedWriter( new FileWriter(newf) );
	
						//output.write(p2+"\n");
						output.write(content);
						//output.write("\n\n\n\n\n\n");
						output.close();
						
					}
				}
				
			}
			//output.close();
			System.out.println("Total Found:"+found+"/"+players.size()+",before duplicated size:"+plist.size());
			
			
		} catch (UnsupportedEncodingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (SQLException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		}
	}
	
	*/
	public void writeIsMarriedTo(String dir){
		try {
			//File dataFile = new File(dir);
			//BufferedWriter output = new BufferedWriter( new FileWriter(dataFile) );
		
			ArrayList<YagoEntity> plist = this.retrieveAll("isMarriedTo");
			IndexTACReader itacr = new IndexTACReader("/largedata1/cmput696/students/lyao1/tacindex");
			
			int found=0;
			HashSet<String> players = new HashSet<String>();
			
			for(YagoEntity e:plist){
				String p1 = e.getArg1();
				
				if( !players.contains(p1)){
					players.add(p1);
					String p2 = itacr.convert(p1);
					
				
					String content = itacr.queryWikiContent(p2);
					if(content!=null){
						
						found++;
						System.out.println("Found:"+p2);
						System.out.println("writing file");
						
						File newf = new File(dir+p2+".txt");
						newf.createNewFile();
						BufferedWriter output = new BufferedWriter( new FileWriter(newf) );
	
						//output.write(p2+"\n");
						output.write(content);
						//output.write("\n\n\n\n\n\n");
						output.close();
						
					}
				}
				
				String p3 = e.getArg2();
				if( !players.contains(p3)){
					players.add(p3);
					String p4 = itacr.convert(p3);
					
				
					String content = itacr.queryWikiContent(p4);
					if(content!=null){
						
						found++;
						System.out.println("Found:"+p4);
						System.out.println("writing file");
						
						File newf = new File(dir+p4+".txt");
						newf.createNewFile();
						BufferedWriter output = new BufferedWriter( new FileWriter(newf) );
	
						//output.write(p2+"\n");
						output.write(content);
						//output.write("\n\n\n\n\n\n");
						output.close();
						
					}
				}
				
			}
			//output.close();
			System.out.println("Total Found:"+found+"/"+players.size()+",before duplicated size:"+plist.size());
			
			
		} catch (UnsupportedEncodingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (SQLException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		}
	}
	
	public static void main(String[]args) throws ClassNotFoundException, SQLException, UnsupportedEncodingException{
		/*Class.forName("org.postgresql.Driver");
		Connection db = DriverManager.getConnection("jdbc:postgresql:yago2spotlx_db", "lyao1", "");
		if( db==null){
			System.out.println("null");
		}else{
			System.out.println("not null");
		}
		Statement st = db.createStatement();
	
		ResultSet rs = st.executeQuery("select arg1,arg2,relation,timebegin from relationalfacts where relation='playsFor' and timebegin<'2010-01-01' limit 10000;");
		PrintStream out = new PrintStream(System.out, true, "UTF-8");
		HashSet<String> players = new HashSet<String>();
		IndexTACReader itacr = new IndexTACReader("/largedata1/cmput696/students/lyao1/tacindex");
		int find_count =0;
		while (rs.next()) {
		    //out.print("Column 1 returned ");
			String p1 = rs.getString(1);
			
			if( !players.contains(p1)){
				players.add(p1);
				String p2 = itacr.convert(p1);
				String content = itacr.queryWikiContent(p2);
				if(content!=null){
					find_count++;
					out.println("Found:"+p2);
				}
			}
			
		   
		}
		rs.close();
		st.close();
		out.println("Total Found:"+find_count+"/"+players.size());
		*/
		ExtractYago ey = new ExtractYago();
		//ey.writePlaysFor("/largedata1/cmput696/students/lyao1/sogodata/1data/texts/");
		ey.writeIsMarriedTo("/largedata1/cmput696/students/lyao1/sogodata/2data/texts/");
		
	}
	
	//get 100 triples then find out their wikitext 
}
